package com.me.study.arithmeticAndDataStructure;

import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * 前缀树查找单词、百度搜索自动提示
 */
public class TrieTreeSearch {

    public static void main(String[] args) {
        String[] goods = new String[]{"huawei phone","huawei router","apple iphone","xiaomi mix"};
        String[] keywords = new String[]{"Huawei", "apple"};

        String[][] result = new TrieTreeSearch().queryGoods(goods, keywords);
        for (String[] strings : result) {
            System.out.println(Arrays.toString(strings));
        }
    }

    public String[][] queryGoods(String[] goods, String[] keywords) {
        Trie trie = new Trie(goods);

        String[][] result = new String[keywords.length][];
        for (int i = 0; i < keywords.length; i++) {
            result[i] = trie.search(keywords[i]);
        }

        return result;
    }
    
    /** 前缀树 */
    public static class Trie {
        private final TrieNode root = new TrieNode();
        private final Map<Integer, String> cache = new HashMap<>();

        public Trie(String[] goods) {
            for (int i = 0; i < goods.length; i++) {
                String good = goods[i];
                cache.put(i, good);
                // 以下标作为商品ID，
                insert(good, i);
            }
        }
        
        /** 插入数据 */
        public void insert(String word, Integer index) {
            TrieNode node = root;
            for (char c : word.toCharArray()) {
                //判断对应节点是否为空， 为空插入
                if (node.next.get(c) == null) {
                    node.next.put(c, new TrieNode(index));
                } else {
                    node.next.get(c).goods.add(index);
                }
                //继续插入下一个
                node = node.next.get(c);
            }
        }

        public String[] search(String prefix) {
            TrieNode node = root;
            for (char c : prefix.toCharArray()) {
                if (node.next.get(c) == null) {
                    return new String[0];
                }
                node = node.next.get(c);
            }

            String[] result = new String[node.goods.size()];
            for (int i = 0; i < node.goods.size(); i++) {
                result[i] = cache.get(node.goods.get(i));
            }

            return result;
        }
    }

    public static class TrieNode {
        private final Map<Character, TrieNode> next;
        private final List<Integer> goods = new LinkedList<>();

        public TrieNode(){
            this.next = new HashMap<>();
        }

        public TrieNode(Integer index){
            this.next = new HashMap<>();
            this.goods.add(index);
        }
    }

}
